import numpy
import os

from PIL import Image
import numpy as np
import random
from keras.models import Model

single_num_path = "/Users/mc/Desktop/captcha_pic_single_num/"  # 要训练的图片目录 6000张 尺寸32x32
single_num_cha_path = "/Users/mc/Desktop/captcha_pic_single_num_cha/"  # 要训练的图片目录 6000张 尺寸32x32
four_num_cha_path = "/Users/mc/Desktop/captcha_pic_4_num_cha/"  # 要训练的图片目录 6000张 尺寸32x32

trainPercentage = 5 / 6  # 训练数据占比 训练5:测试1

# type  1为1位数字  2为1位数字或字母 3为4位数字和字母组合
def loadData(type):
    if type == '1':
        path = single_num_path
    elif type == '2':
        path = single_num_cha_path
    elif type == '3':
        path = four_num_cha_path

    pics = os.listdir(path)
    pics.remove(".DS_Store")  # 删除mac文件夹下的隐藏文件
    random.shuffle(pics)  # 打乱顺序
    picSize = len(pics)
    trainSize = int(picSize * trainPercentage)

    x_train, y_train = getData(pics, 0, trainSize, path)
    x_test, y_test = getData(pics, trainSize, picSize, path)
    return (x_train, y_train), (x_test, y_test)


def getData(pics, start, end, path):
    x_train = []
    y_train = []
    for i in range(start, end):
        file = pics[i]
        if file[-3:] == "png":
            x_train.append(numpy.asarray(Image.open(path + file)))
            y_asc = ord(file[:1])
            if (47< y_asc< 58):
                y = y_asc -48
            elif(96 < y_asc <123):
                y = y_asc - 87
            y_train.append(numpy.asarray([y]))
    return np.asarray(x_train), np.asarray(y_train)


# loadData('1')
